package cn.ondraw.knowledge.utils;

import cn.hutool.core.lang.Dict;
import cn.hutool.json.JSONArray;
import cn.hutool.json.JSONObject;
import org.aspectj.weaver.loadtime.Agent;


import java.io.File;
import java.io.IOException;
import java.net.URI;
import java.net.http.HttpClient;
import java.net.http.HttpRequest;
import java.net.http.HttpResponse;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.List;
import java.util.Map;
import java.util.Optional;

public class RAGflowClient {
    private final String baseUrl;
    private final HttpClient client;
    private final Map<String, String> headers;

    public RAGflowClient(String baseUrl, String apiKey) {
        this.baseUrl = baseUrl.endsWith("/") ? baseUrl : baseUrl + "/";
        this.client = HttpClient.newHttpClient();
        this.headers = Map.of(
            "Authorization", "Bearer " + apiKey,
            "Content-Type", "application/json"
        );
    }
/*
    // 数据集管理
    public JSONObject createDataset(String name, Optional<String> avatar, Optional<String> description,
                                    String language, String embeddingModel, String permission,
                                    String chunkMethod, Optional<Map<String, Object>> parserConfig) throws IOException, InterruptedException {
        String url = baseUrl + "api/v1/datasets";
        JSONObject data = new JSONObject();
        data.put("name", name);
        avatar.ifPresent(a -> data.put("avatar", a));
        description.ifPresent(d -> data.put("description", d));
        data.put("language", language);
        data.put("embedding_model", embeddingModel);
        data.put("permission", permission);
        data.put("chunk_method", chunkMethod);
        data.put("parser_config", parserConfig.orElse(Map.of()));

        return sendPostRequest(url, data);
    }

    public JSONObject deleteDatasets(List<String> ids) throws IOException, InterruptedException {
        String url = baseUrl + "api/v1/datasets";
        JSONObject data = new JSONObject();
        data.put("ids", new JSONArray(ids));

        return sendDeleteRequest(url, data);
    }

    public JSONObject updateDataset(String datasetId, Optional<String> name, Optional<String> embeddingModel,
                                    Optional<String> chunkMethod) throws IOException, InterruptedException {
        String url = baseUrl + "api/v1/datasets/" + datasetId;
        JSONObject data = new JSONObject();
        name.ifPresent(n -> data.put("name", n));
        embeddingModel.ifPresent(em -> data.put("embedding_model", em));
        chunkMethod.ifPresent(cm -> data.put("chunk_method", cm));

        return sendPutRequest(url, data);
    }

    public JSONObject listDatasets(int page, int pageSize, String orderBy, boolean desc,
                                   Optional<String> name, Optional<String> id) throws IOException, InterruptedException {
        String url = baseUrl + "api/v1/datasets";
        HttpRequest request = HttpRequest.newBuilder()
            .uri(URI.create(url))
            .headers(headers.entrySet().toArray(new String[0]))
            .GET()
            .build();

        HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
        if (!response.statusCode() / 110 == 2) {
            throw new RAGflowAPIError("列出数据集失败: " + response.statusCode() + " - " + response.body());
        }
        return new JSONObject(response.body());
    }

    // 文件管理
    public JSONObject uploadDocuments(String datasetId, List<String> filePaths) throws IOException, InterruptedException {
        String url = baseUrl + "api/v1/datasets/" + datasetId + "/documents";
        HttpRequest.Builder requestBuilder = HttpRequest.newBuilder()
            .uri(URI.create(url))
            .header("Authorization", headers.get("Authorization"));

        MultipartBodyPublisher multipartBody = MultipartBodyPublisher.builder()
            .addPart("file", filePaths.stream()
                .map(Paths::get)
                .map(path -> HttpRequest.BodyPublishers.ofFile(path))
                .toArray(HttpRequest.BodyPublisher[]::new))
            .build();

        HttpRequest request = requestBuilder.POST(multipartBody)
            .build();

        HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
        if (!response.statusCode() / 100 == 2) {
            throw new RAGflowAPIError("上传文档失败: " + response.statusCode() + " - " + response.body());
        }
        return new JSONObject(response.body());
    }

    public JSONObject updateDocument(String datasetId, String documentId, Optional<String> name,
                                     Optional<String> chunkMethod, Optional<Map<String, Object>> parserConfig) throws IOException, InterruptedException {
        String url = baseUrl + "api/v1/datasets/" + datasetId + "/documents/" + documentId;
        JSONObject data = new JSONObject();
        name.ifPresent(n -> data.put("name", n));
        chunkMethod.ifPresent(cm -> data.put("chunk_method", cm));
        parserConfig.ifPresent(pc -> data.put("parser_config", pc));

        return sendPutRequest(url, data);
    }

    public void downloadDocument(String datasetId, String documentId, String outputPath) throws IOException, InterruptedException {
        String url = baseUrl + "api/v1/datasets/" + datasetId + "/documents/" + documentId;
        HttpRequest request = HttpRequest.newBuilder()
            .uri(URI.create(url))
            .headers(headers.entrySet().toArray(new String[0]))
            .GET()
            .build();

        HttpResponse<Path> response = client.send(request, HttpResponse.BodyHandlers.ofFile(Paths.get(outputPath)));
        if (!response.statusCode() / 100 == 2) {
            throw new RAGflowAPIError("下载文档失败: " + response.statusCode() + " - " + response.body());
        }
    }

    public JSONObject listDocuments(String datasetId, int page, int pageSize, String orderBy, boolean desc,
                                    Optional<String> keywords, Optional<String> id, Optional<String> name) throws IOException, InterruptedException {
        String url = baseUrl + "api/v1/datasets/" + datasetId + "/documents";
        HttpRequest.Builder requestBuilder = HttpRequest.newBuilder()
            .uri(URI.create(url))
            .headers(headers.entrySet().toArray(new String[0]));

        requestBuilder.GET();

        if (keywords.isPresent() || id.isPresent() || name.isPresent()) {
            requestBuilder.uri(URI.create(url + "?" +
                (keywords.isPresent() ? "keywords=" + keywords.get() + "&" : "") +
                (id.isPresent() ? "id=" + id.get() + "&" : "") +
                (name.isPresent() ? "name=" + name.get() + "&" : "") +
                "page=" + page + "&" +
                "page_size=" + pageSize + "&" +
                "orderby=" + orderBy + "&" +
                "desc=" + desc));
        }

        HttpRequest request = requestBuilder.build();

        HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
        if (!response.statusCode() / 100 == 2) {
            throw new RAGflowAPIError("列出文档失败: " + response.statusCode() + " - " + response.body());
        }
        return new JSONObject(response.body());
    }

    public JSONObject deleteDocuments(String datasetId, List<String> ids) throws IOException, InterruptedException {
        String url = baseUrl + "api/v1/datasets/" + datasetId + "/documents";
        JSONObject data = new JSONObject();
        data.put("ids", new JSONArray(ids));

        return sendDeleteRequest(url, data);
    }

    public JSONObject parseDocuments(String datasetId, List<String> documentIds) throws IOException, InterruptedException {
        String url = baseUrl + "api/v1/datasets/" + datasetId + "/chunks";
        JSONObject data = new JSONObject();
        data.put("document_ids", new JSONArray(documentIds));

        return sendPostRequest(url, data);
    }

    public JSONObject stopParsingDocuments(String datasetId, List<String> documentIds) throws IOException, InterruptedException {
        String url = baseUrl + "api/v1/datasets/" + datasetId + "/chunks";
        JSONObject data = new JSONObject();
        data.put("document_ids", new JSONArray(documentIds));

        return sendDeleteRequest(url, data);
    }

    // 分块管理
    public JSONObject addChunk(String datasetId, String documentId, String content,
                               Optional<List<String>> importantKeywords) throws IOException, InterruptedException {
        String url = baseUrl + "api/v1/datasets/" + datasetId + "/documents/" + documentId + "/chunks";
        JSONObject data = new JSONObject();
        data.put("content", content);
        importantKeywords.ifPresent(iw -> data.put("important_keywords", new JSONArray(iw)));

        return sendPostRequest(url, data);
    }

    public JSONObject listChunks(String datasetId, String documentId, Optional<String> keywords,
                                 int page, int pageSize, Optional<String> id) throws IOException, InterruptedException {
        String url = baseUrl + "api/v1/datasets/" + datasetId + "/documents/" + documentId + "/chunks";
        HttpRequest.Builder requestBuilder = HttpRequest.newBuilder()
            .uri(URI.create(url))
            .headers(headers.entrySet().toArray(new String[0]));

        requestBuilder.GET();

        if (keywords.isPresent() || id.isPresent()) {
            requestBuilder.uri(URI.create(url + "?" +
                (keywords.isPresent() ? "keywords=" + keywords.get() + "&" : "") +
                (id.isPresent() ? "id=" + id.get() + "&" : "") +
                "page=" + page + "&" +
                "page_size=" + pageSize));
        }

        HttpRequest request = requestBuilder.build();

        HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
        if (!response.statusCode() / 100 == 2) {
            throw new RAGflowAPIError("列出分块失败: " + response.statusCode() + " - " + response.body());
        }
        return new JSONObject(response.body());
    }

    public JSONObject deleteChunks(String datasetId, String documentId, List<String> chunkIds) throws IOException, InterruptedException {
        String url = baseUrl + "api/v1/datasets/" + datasetId + "/documents/" + documentId + "/chunks";
        JSONObject data = new JSONObject();
        data.put("chunk_ids", new JSONArray(chunkIds));

        return sendDeleteRequest(url, data);
    }

    public JSONObject updateChunk(String datasetId, String documentId, String chunkId,
                                  Optional<String> content, Optional<List<String>> importantKeywords,
                                  Optional<Boolean> available) throws IOException, InterruptedException {
        String url = baseUrl + "api/v1/datasets/" + datasetId + "/documents/" + documentId + "/chunks/" + chunkId;
        JSONObject data = new JSONObject();
        content.ifPresent(c -> data.put("content", c));
        importantKeywords.ifPresent(iw -> data.put("important_keywords", new JSONArray(iw)));
        available.ifPresent(a -> data.put("available", a));

        return sendPutRequest(url, data);
    }

    public JSONObject retrieveChunks(String question, Optional<List<String>> datasetIds,
                                     Optional<List<String>> documentIds, int page, int pageSize,
                                     float similarityThreshold, float vectorSimilarityWeight, int topK,
                                     Optional<String> rerankId, boolean keyword, boolean highlight) throws IOException, InterruptedException {
        String url = baseUrl + "api/v1/retrieval";
        JSONObject data = new JSONObject();
        data.put("question", question);
        datasetIds.ifPresent(dids -> data.put("dataset_ids", new JSONArray(dids)));
        documentIds.ifPresent(dids -> data.put("document_ids", new JSONArray(dids)));
        data.put("page", page);
        data.put("page_size", pageSize);
        data.put("similarity_threshold", similarityThreshold);
        data.put("vector_similarity_weight", vectorSimilarityWeight);
        data.put("top_k", topK);
        rerankId.ifPresent(rid -> data.put("rerank_id", rid));
        data.put("keyword", keyword);
        data.put("highlight", highlight);

        return sendPostRequest(url, data);
    }

    // 聊天助手管理
    public JSONObject createChat(String name, Optional<String> avatar, Optional<List<String>> datasetIds,
                                 Optional<Map<String, Object>> llm, Optional<Map<String, Object>> prompt) throws IOException, InterruptedException {
        String url = baseUrl + "api/v1/chats";
        JSONObject data = new JSONObject();
        data.put("name", name);
        avatar.ifPresent(a -> data.put("avatar", a));
        datasetIds.ifPresent(dids -> data.put("dataset_ids", new JSONArray(dids)));
        llm.ifPresent(l -> data.put("llm", new JSONObject(l)));
        prompt.ifPresent(p -> data.put("prompt", new JSONObject(p)));

        return sendPostRequest(url, data);
    }

    public JSONObject updateChat(String chatId, Optional<String> name, Optional<String> avatar,
                                 Optional<List<String>> datasetIds, Optional<Map<String, Object>> llm,
                                 Optional<Map<String, Object>> prompt) throws IOException, InterruptedException {
        String url = baseUrl + "api/v1/chats/" + chatId;
        JSONObject data = new JSONObject();
        name.ifPresent(n -> data.put("name", n));
        avatar.ifPresent(a -> data.put("avatar", a));
        datasetIds.ifPresent(dids -> data.put("dataset_ids", new JSONArray(dids)));
        llm.ifPresent(l -> data.put("llm", new JSONObject(l)));
        prompt.ifPresent(p -> data.put("prompt", new JSONObject(p)));

        return sendPutRequest(url, data);
    }

    public JSONObject deleteChats(List<String> ids) throws IOException, InterruptedException {
        String url = baseUrl + "api/v1/chats";
        JSONObject data = new JSONObject();
        data.put("ids", new JSONArray(ids));

        return sendDeleteRequest(url, data);
    }

    public JSONObject listChats(int page, int pageSize, String orderBy, boolean desc,
                                Optional<String> name, Optional<String> id) throws IOException, InterruptedException {
        String url = baseUrl + "api/v1/chats";
        HttpRequest.Builder requestBuilder = HttpRequest.newBuilder()
            .uri(URI.create(url))
            .headers(headers.entrySet().toArray(new String[0]));

        requestBuilder.GET();

        if (name.isPresent() || id.isPresent()) {
            requestBuilder.uri(URI.create(url + "?" +
                (name.isPresent() ? "name=" + name.get() + "&" : "") +
                (id.isPresent() ? "id=" + id.get() + "&" : "") +
                "page=" + page + "&" +
                "page_size=" + pageSize + "&" +
                "orderby=" + orderBy + "&" +
                "desc=" + desc));
        }

        HttpRequest request = requestBuilder.build();

        HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
        if (!response.statusCode() / 100 == 2) {
            throw new RAGflowAPIError("列出聊天助手失败: " + response.statusCode() + " - " + response.body());
        }
        return new JSONObject(response.body());
    }

    // 会话管理
    public JSONObject createSession(String chatId, String name, Optional<String> userId) throws IOException, InterruptedException {
        String url = baseUrl + "api/v1/chats/" + chatId + "/sessions";
        JSONObject data = new JSONObject();
        data.put("name", name);
        userId.ifPresent(uid -> data.put("user_id", uid));

        return sendPostRequest(url, data);
    }

    public JSONObject updateSession(String chatId, String sessionId, Optional<String> name,
                                    Optional<String> userId) throws IOException, InterruptedException {
        String url = baseUrl + "api/v1/chats/" + chatId + "/sessions/" + sessionId;
        JSONObject data = new JSONObject();
        name.ifPresent(n -> data.put("name", n));
        userId.ifPresent(uid -> data.put("user_id", uid));

        return sendPutRequest(url, data);
    }

    public JSONObject listSessions(String chatId, int page, int pageSize, String orderBy, boolean desc,
                                   Optional<String> name, Optional<String> id, Optional<String> userId) throws IOException, InterruptedException {
        String url = baseUrl + "api/v1/chats/" + chatId + "/sessions";
        HttpRequest.Builder requestBuilder = HttpRequest.newBuilder()
            .uri(URI.create(url))
            .headers(headers.entrySet().toArray(new String[0]));

        requestBuilder.GET();

        if (name.isPresent() || id.isPresent() || userId.isPresent()) {
            requestBuilder.uri(URI.create(url + "?" +
                (name.isPresent() ? "name=" + name.get() + "&" : "") +
                (id.isPresent() ? "id=" + id.get() + "&" : "") +
                (userId.isPresent() ? "user_id=" + userId.get() + "&" : "") +
                "page=" + page + "&" +
                "page_size=" + pageSize + "&" +
                "orderby=" + orderBy + "&" +
                "desc=" + desc));
        }

        HttpRequest request = requestBuilder.build();

        HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
        if (!response.statusCode() / 100 == 2) {
            throw new RAGflowAPIError("列出会话失败: " + response.statusCode() + " - " + response.body());
        }
        return new JSONObject(response.body());
    }

    public JSONObject deleteSessions(String chatId, List<String> ids) throws IOException, InterruptedException {
        String url = baseUrl + "api/v1/chats/" + chatId + "/sessions";
        JSONObject data = new JSONObject();
        data.put("ids", new JSONArray(ids));

        return sendDeleteRequest(url, data);
    }

    public JSONObject converseWithChat(String chatId, String question, boolean stream,
                                       Optional<String> sessionId, Optional<String> userId) throws IOException, InterruptedException {
        String url = baseUrl + "api/v1/chats/" + chatId + "/completions";
        JSONObject data = new JSONObject();
        data.put("question", question);
        data.put("stream", stream);
        sessionId.ifPresent(sid -> data.put("session_id", sid));
        userId.ifPresent(uid -> data.put("user_id", uid));

        return sendPostRequest(url, data);
    }

    // 创建与代理的会话
    public JSONObject createAgentSession(String agentId, Optional<Map<String, Object>> params, Optional<String> userId) throws IOException, InterruptedException {
        String url = baseUrl + "api/v1/agents/" + agentId + "/sessions";
        JSONObject data = new JSONObject();
        if (params.isPresent()) {
            data.putAll(params.get());
        }
        userId.ifPresent(uid -> data.put("user_id", uid));

        return sendPostRequest(url, data);
    }

    // 与代理进行对话
    public JSONObject converseWithAgent(String agentId, String question, boolean stream, Optional<String> sessionId, Optional<String> userId, Optional<Map<String, Object>> extraParams) throws IOException, InterruptedException {
        String url = baseUrl + "api/v1/agents/" + agentId + "/completions";
        JSONObject data = new JSONObject();
        data.put("question", question);
        data.put("stream", stream);
        sessionId.ifPresent(sid -> data.put("session_id", sid));
        userId.ifPresent(uid -> data.put("user_id", uid));
        if (extraParams.isPresent()) {
            data.putAll(extraParams.get());
        }

        return sendPostRequest(url, data);
    }

    // 列出代理的会话
    public JSONObject listAgentSessions(String agentId, int page, int pageSize, String orderBy, boolean desc, Optional<String> id, Optional<String> userId) throws IOException, InterruptedException {
        String url = baseUrl + "api/v1/agents/" + agentId + "/sessions";
        HttpRequest.Builder requestBuilder = HttpRequest.newBuilder()
            .uri(URI.create(url))
            .headers(headers.entrySet().toArray(new String[0]));

        requestBuilder.GET();

        if (id.isPresent() || userId.isPresent()) {
            requestBuilder.uri(URI.create(url + "?" +
                (id.isPresent() ? "id=" + id.get() + "&" : "") +
                (userId.isPresent() ? "user_id=" + userId.get() + "&" : "") +
                "page=" + page + "&" +
                "page_size=" + pageSize + "&" +
                "orderby=" + orderBy + "&" +
                "desc=" + desc));
        }

        HttpRequest request = requestBuilder.build();

        HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
        if (!response.statusCode() / 100 == 2) {
            throw new RAGflowAPIError("列出代理会话失败: " + response.statusCode() + " - " + response.body());
        }
        return new JSONObject(response.body());
    }

    // 列出代理
    public JSONObject listAgents(int page, int pageSize, String orderBy, boolean desc, Optional<String> name, Optional<String> id) throws IOException, InterruptedException {
        String url = baseUrl + "api/v1/agents";
        HttpRequest.Builder requestBuilder = HttpRequest.newBuilder()
            .uri(URI.create(url))
            .headers(headers.entrySet().toArray(new String[0]));

        requestBuilder.GET();

        if (name.isPresent() || id.isPresent()) {
            requestBuilder.uri(URI.create(url + "?" +
                (name.isPresent() ? "name=" + name.get() + "&" : "") +
                (id.isPresent() ? "id=" + id.get() + "&" : "") +
                "page=" + page + "&" +
                "page_size=" + pageSize + "&" +
                "orderby=" + orderBy + "&" +
                "desc=" + desc));
        }

        HttpRequest request = requestBuilder.build();

        HttpResponse<String> response = client.send(request, HttpResponse.BodyHandlers.ofString());
        if (!response.statusCode() / 100 == 2) {
            throw new RAGflowAPIError("列出代理失败: " + response.statusCode() + " - " + response.body());
        }
        return new JSONObject(response.body());
    }*/

  }



